import pygame
import serial
import time
import keyboard
###############
#   function description:
#       this script makes joystick the controler of the FPV, the controling logic is as follows:
#           left stick: control output power
#           right sitck and right buttons: control PID parameters of pitch and roll 
#               right stick controls pitch or roll PID paras increse or decrease
#               right buttons controls PID para, 2 to P, 3 to I, 1 to D.
###############
# 初始化手柄
pygame.init()
pygame.joystick.init()
joystick = pygame.joystick.Joystick(0)
joystick.init()
# 配置蓝牙串口
bluetooth_port = "COM5"  # 替换为你的串口号
baudrate = 115200         # 需与 ESP32 的波特率一致

# global variables
time_delay = 0.5
# combine buttons on joystick with PID para
# BUTTONS
PID_P  = 2
PID_I = 3
PID_D = 1
START = 8
PAUSE = 10
PRINT = 4
CHANGE_PID_MODE = 0 # 改变pid改变模式，在外环模式和内环模式切换
# AXIS
PITCH  = 3
ROLL = 2
YAW_UP = 4
YAW_DOWN= 5
# pid adjust step
AJUST_STEP = 1 

# command list
COMMAND_LIST = [["pp+","pp-"],
                ["pr+","pr-"],
                ["py+","py-"],
                ["ip+","ip-"],
                ["ir+","ir-"],
                ["iy+","iy-"],
                ["dp+","dp-"],
                ["dr+","dr-"],
                ["dy+","dy-"]
                ]
# 初始化蓝牙串口
# bluetooth = serial.Serial('COM6', 9600, timeout=1)  # 修改为你的串口

print("手柄控制已启动！摇杆上下推发送指令...")
# 函数
def precise_delay(delay_seconds):
    start = time.perf_counter()
    while time.perf_counter() - start < delay_seconds:
        pass  # 忙等待（占用 CPU）

class ypr_pid:
    def __init__(self):
        self.pitch  =   [5,0,0]
        self.roll   =   [5,0,0]
        self.yaw    =   [5,0,0]
        self.pitch_w  =   [30,0,2]
        self.roll_w   =   [30,0,2]
        self.yaw_w    =   [10,1,1]

    def change_pitch(self,pid,up_down):
        if(up_down == 0):
            self.pitch[pid] = self.pitch[pid] - AJUST_STEP
            return self.pitch[pid]
        elif(up_down == 1):
            self.pitch[pid] = self.pitch[pid] + AJUST_STEP
            return self.pitch[pid]
    def change_roll(self,pid,up_down):
        if(up_down == 0):
            self.roll[pid] = self.roll[pid] - AJUST_STEP
            return self.roll[pid]
        elif(up_down == 1):
            self.roll[pid] = self.roll[pid] + AJUST_STEP
            return self.roll[pid]

    def change_yaw(self,pid,up_down):
        if(up_down == 0):
            self.yaw[pid] = self.yaw[pid] - AJUST_STEP
            return self.yaw[pid]
        elif(up_down == 1):
            self.yaw[pid] = self.yaw[pid] + AJUST_STEP
            return self.yaw[pid]
         
    def change_pitch_w(self,pid,up_down):
        if(up_down == 0):
            self.pitch_w[pid] = self.pitch_w[pid] - AJUST_STEP
            return self.pitch_w[pid]
        elif(up_down == 1):
            self.pitch_w[pid] = self.pitch_w[pid] + AJUST_STEP
            return self.pitch_w[pid]
        
    def change_roll_w(self,pid,up_down):
        if(up_down == 0):
            self.roll_w[pid] = self.roll_w[pid] - AJUST_STEP
            return self.roll_w[pid]
        elif(up_down == 1):
            self.roll_w[pid] = self.roll_w[pid] + AJUST_STEP
            return self.roll_w[pid]
    
    def change_yaw_w(self,pid,up_down):
        if(up_down == 0):
            self.yaw_w[pid] = self.yaw_w[pid] - AJUST_STEP
            return self.yaw_w[pid]
        elif(up_down == 1):
            self.yaw_w[pid] = self.yaw_w[pid] + AJUST_STEP
            return self.yaw_w[pid]

def change_pid(command:str, ypr_pid:ypr_pid):
    for i,J in enumerate(COMMAND_LIST):
        if(J[0] == command):
            if(i%3 == 0):
                if not is_mode_inner: 
                    val = ypr_pid.change_pitch(i//3,1)
                elif is_mode_inner:
                    val = ypr_pid.change_pitch_w(i//3,1)
                break
            elif(i%3 == 1):
                if not is_mode_inner:
                    val = ypr_pid.change_roll(i//3,1)
                elif is_mode_inner:
                    val = ypr_pid.change_roll_w(i//3,1)
                break
            elif(i%3 == 2):
                if not is_mode_inner:
                    val =ypr_pid.change_yaw(i//3,1)
                elif is_mode_inner:
                    val = ypr_pid.change_yaw_w(i//3,1)
                break
        elif(J[1] == command):
            if(i%3 == 0):
                if not is_mode_inner: 
                    val = ypr_pid.change_pitch(i//3,0)
                elif is_mode_inner:
                    val = ypr_pid.change_pitch_w(i//3,0)
                break
            elif(i%3 == 1):
                if not is_mode_inner:
                    val = ypr_pid.change_roll(i//3,0)
                elif is_mode_inner:
                    val = ypr_pid.change_roll_w(i//3,0)
                break
            elif(i%3 == 2):
                if not is_mode_inner:
                    val =ypr_pid.change_yaw(i//3,0)
                elif is_mode_inner:
                    val = ypr_pid.change_yaw_w(i//3,0)
                break
    return command[0:2] + str(val) + '\n'

if __name__ == "__main__":
    # create bluetooth object
    bluetooth = serial.Serial(bluetooth_port, baudrate, timeout=1)
    print(f"已连接蓝牙串口: {bluetooth_port}")
    # create ypr_pid object
    my_ypr_pid = ypr_pid()
    # running state
    running = True
    is_mode_inner = False 
    # pid change state
    while running:
        # 检测键盘 Q 键按下
        pygame.event.pump()

        # 读取蓝牙指令
        if bluetooth.in_waiting > 0:
            data = bluetooth.readline().decode('utf-8').strip()  # 读取一行并解码为字符串
            print(f"收到数据: {data}")

        # 退出,检测键盘是否按下k, short for kill
        if keyboard.is_pressed('k'):
            print('The k key is being pressed')
            running = False  

        # 向esp发送s，代表开始，short for start
        if joystick.get_button(START):
            # print('The \"q\q" key is being pressed')
            message = "s\n"
            bluetooth.write(message.encode('utf-8'))
            precise_delay(time_delay)

        # 向esp发送q,代表停止，short for quit
        if joystick.get_button(PAUSE):
            # print('The \"q\q" key is being pressed')
            message = "q\n"
            bluetooth.write(message.encode('utf-8'))
            precise_delay(time_delay)

        # 发送p， 打印PWM 
        if joystick.get_button(PRINT):
            # print('The \"q\q" key is being pressed')
            message = "p\n"
            bluetooth.write(message.encode('utf-8'))
            precise_delay(time_delay)

        # 发送m， 改变pid调整模式，在外环和内环切换 
        if joystick.get_button(CHANGE_PID_MODE):
            # print('The \"q\q" key is being pressed')
            message = "m\n"
            is_mode_inner = not is_mode_inner
            print(f"in com, change mode is {is_mode_inner}")
            bluetooth.write(message.encode('utf-8'))
            precise_delay(time_delay)

        ## 轮询所有按钮状态
        # for button in range(joystick.get_numbuttons()):
        #     if joystick.get_button(button):
        #         print(f"按钮 {button} 被按下")

        # pid for pitch and roll
        # pp+-
        if joystick.get_button(PID_P) and abs(joystick.get_axis(PITCH)) > 0.5:
            if joystick.get_axis(PITCH) < 0: 
                message = "pp+"
                print(message )
                message = change_pid(message, my_ypr_pid)
                bluetooth.write(message.encode('utf-8'))
                # print(message)
                precise_delay(time_delay )
            elif joystick.get_axis(PITCH) > 0: 
                message = "pp-"
                print(message )
                message = change_pid(message, my_ypr_pid)
                bluetooth.write(message.encode('utf-8'))
                # print(message)
                precise_delay(time_delay )
        # pr+-
        elif joystick.get_button(PID_P) and abs(joystick.get_axis(ROLL)) > 0.5:
            if joystick.get_axis(ROLL) < 0: 
                message = "pr+"
                print(message )
                message = change_pid(message, my_ypr_pid)
                bluetooth.write(message.encode('utf-8'))
                precise_delay(time_delay )
            elif joystick.get_axis(ROLL) > 0: 
                message = "pr-"
                print(message )
                message = change_pid(message, my_ypr_pid)
                bluetooth.write(message.encode('utf-8'))
                precise_delay(time_delay )
        # ip +-
        elif joystick.get_button(PID_I) and abs(joystick.get_axis(PITCH)) > 0.5:
            if joystick.get_axis(PITCH) < 0: 
                message = "ip+"
                print(message )
                message = change_pid(message, my_ypr_pid)
                bluetooth.write(message.encode('utf-8'))
                precise_delay(time_delay )
            elif joystick.get_axis(PITCH) > 0: 
                message = "ip-"
                print(message )
                message = change_pid(message, my_ypr_pid)
                bluetooth.write(message.encode('utf-8'))
                precise_delay(time_delay )
        # ir +-
        elif joystick.get_button(PID_I) and abs(joystick.get_axis(ROLL)) > 0.5:
            if joystick.get_axis(ROLL) < 0: 
                message = "ir+"
                print(message )
                message = change_pid(message, my_ypr_pid)
                bluetooth.write(message.encode('utf-8'))
                precise_delay(time_delay )
            elif joystick.get_axis(ROLL) > 0: 
                message = "ir-"
                print(message )
                message = change_pid(message, my_ypr_pid)
                bluetooth.write(message.encode('utf-8'))
                precise_delay(time_delay )
        # dp +-
        elif joystick.get_button(PID_D) and abs(joystick.get_axis(PITCH)) > 0.5:
            if joystick.get_axis(PITCH) < 0: 
                message = "dp+"
                print(message )
                message = change_pid(message, my_ypr_pid)
                bluetooth.write(message.encode('utf-8'))
                precise_delay(time_delay )
            elif joystick.get_axis(PITCH) > 0: 
                message = "dp-"
                print(message )
                message = change_pid(message, my_ypr_pid)
                bluetooth.write(message.encode('utf-8'))
                precise_delay(time_delay )
        # dr +-
        elif joystick.get_button(PID_D) and abs(joystick.get_axis(ROLL)) > 0.5:
            if joystick.get_axis(ROLL) < 0: 
                message = "dr+"
                print(message )
                message = change_pid(message, my_ypr_pid)
                bluetooth.write(message.encode('utf-8'))
                precise_delay(time_delay )
            elif joystick.get_axis(ROLL) > 0: 
                message = "dr-"
                print(message )
                message = change_pid(message, my_ypr_pid)
                bluetooth.write(message.encode('utf-8'))
                precise_delay(time_delay )
        # py+-
        elif joystick.get_button(PID_P) and (abs(joystick.get_axis(YAW_UP)- -1) > 0.5 or abs(joystick.get_axis(YAW_DOWN) - -1)) > 0.5:
            if abs(joystick.get_axis(YAW_UP)- -1) > 0.5:
                message = "py+"
                print(message )
                message = change_pid(message, my_ypr_pid)
                bluetooth.write(message.encode('utf-8'))
                precise_delay(time_delay )
            elif abs(joystick.get_axis(YAW_DOWN)- -1) > 0.5:
                message = "py-"
                print(message )
                message = change_pid(message, my_ypr_pid)
                bluetooth.write(message.encode('utf-8'))
                precise_delay(time_delay )
        # iy+-
        elif joystick.get_button(PID_I) and (abs(joystick.get_axis(YAW_UP)- -1) > 0.5 or abs(joystick.get_axis(YAW_DOWN)- -1)) > 0.5:
            if abs(joystick.get_axis(YAW_UP)- -1) > 0.5:
                message = "iy+"
                print(message )
                message = change_pid(message, my_ypr_pid)
                bluetooth.write(message.encode('utf-8'))
                precise_delay(time_delay )
            elif abs(joystick.get_axis(YAW_DOWN)- -1) > 0.5:
                message = "iy-"
                print(message )
                message = change_pid(message, my_ypr_pid)
                bluetooth.write(message.encode('utf-8'))
                precise_delay(time_delay )
        # dy+-
        elif joystick.get_button(PID_D) and (abs(joystick.get_axis(YAW_UP)- -1) > 0.5 or abs(joystick.get_axis(YAW_DOWN)- -1)) > 0.5:
            if abs(joystick.get_axis(YAW_UP)- -1) > 0.5:
                message = "dy+"
                print(message )
                message = change_pid(message, my_ypr_pid)
                bluetooth.write(message.encode('utf-8'))
                precise_delay(time_delay )
            elif abs(joystick.get_axis(YAW_DOWN)- -1) > 0.5:
                message = "dy-"
                print(message )
                message = change_pid(message, my_ypr_pid)
                bluetooth.write(message.encode('utf-8'))
                precise_delay(time_delay )


        # xy 平面的移动
        elif abs(joystick.get_axis(PITCH)) > 0.5: 
            if joystick.get_axis(PITCH) < 0:
                message = "xf\n"
                print(message )
                bluetooth.write(message.encode('utf-8'))
                precise_delay(time_delay )
            elif joystick.get_axis(PITCH) > 0:
                message = "xb\n"
                print(message )
                bluetooth.write(message.encode('utf-8'))
                precise_delay(time_delay )
        elif abs(joystick.get_axis(ROLL)) > 0.5: 
            if joystick.get_axis(ROLL) < 0:
                message = "yf\n"
                print(message )
                bluetooth.write(message.encode('utf-8'))
                precise_delay(time_delay )
            elif joystick.get_axis(ROLL) > 0:
                message = "yb\n"
                print(message )
                bluetooth.write(message.encode('utf-8'))
                precise_delay(time_delay )

        # 轮询所有摇杆轴的值（范围: -1.0 到 1.0）
        for axis in range(0,2):
            axis_value = joystick.get_axis(axis)
            # 左摇杆，竖直推，表示上升下降
            if axis == 1 and abs(axis_value) > 0.5:  # 忽略微小波动
                print(f"轴 {axis} 的值: {axis_value:.2f}")
                if axis_value < 0: # 上推
                    message = "u\n"
                    bluetooth.write(message.encode('utf-8'))
                elif axis_value > 0: #下拉 
                    message = "d\n"
                    bluetooth.write(message.encode('utf-8'))
                precise_delay(time_delay )
            # 左摇杆，横向推，表示左转右转
            elif axis == 0 and abs(axis_value) > 0.5:
                print(f"轴 {axis} 的值: {axis_value:.2f}")
                if axis_value > 0: # 右推
                    message = "r\n"
                    bluetooth.write(message.encode('utf-8'))
                elif axis_value < 0: # 左推
                    message = "l\n"
                    bluetooth.write(message.encode('utf-8'))
                precise_delay(time_delay )

        # change pid for yaw
        # 轮询方向键（Hat）状态（返回 (x, y) 元组，值: -1, 0, 1）
        # for hat in range(joystick.get_numhats()):
        #     hat_value = joystick.get_hat(hat)
        #     if hat_value != (0, 0):
        #         print(f"方向键 {hat} 的值: {hat_value}")
        #         # 上键
        #         if hat_value == (0,1):
        #             message = "py+"
        #             print(message )
        #             message = change_pid(message, my_ypr_pid)
        #             bluetooth.write(message.encode('utf-8'))
        #             precise_delay(time_delay )
        #         # 下键
        #         elif hat_value == (0,-1):
        #             message = "py-"
        #             print(message )
        #             message = change_pid(message, my_ypr_pid)
        #             bluetooth.write(message.encode('utf-8'))
        #             precise_delay(time_delay )
        #         # 左键
        #         elif hat_value == (-1,0):
        #             message = "iy+"
        #             print(message )
        #             message = change_pid(message, my_ypr_pid)
        #             bluetooth.write(message.encode('utf-8'))
        #             precise_delay(time_delay )
        #         # 右键
        #         elif hat_value == (1,0):
        #             message = "iy-"
        #             print(message )
        #             message = change_pid(message, my_ypr_pid)
        #             bluetooth.write(message.encode('utf-8'))
        #             precise_delay(time_delay )
        #         precise_delay(time_delay )

# # 主循环
# if __name__ == "__main__":        
#     running = True
#     while running:
#         # 检测键盘 Q 键按下
#         if keyboard.is_pressed('q'):
#             print('The \"q\q" key is being pressed')
#             running = False     
#         event_list = pygame.event.get()
#         for i,event in enumerate(event_list):
#             if event.type == pygame.QUIT:  # 点击窗口关闭按钮
#                 running = False       
#             # 手柄按钮按下事件
#             elif event.type == pygame.JOYBUTTONDOWN:
#                 print(f"按钮 {event.button} 按下")
#                 precise_delay(0.5)

#             # 摇杆移动事件
#             elif event.type == pygame.JOYAXISMOTION:
#                 if event.axis == 1 and abs(event.value) > 0.5:  # 忽略微小波动
#                     print(f"摇杆 {event.axis} 的值: {event.value:.2f}")
#                     precise_delay(0.5)
#                     # 移动一次摇杆，触发一次
#                     pygame.event.clear()
#                     break
#                 elif event.axis == 3 and abs(event.value) > 0.5:
#                     if joystick.get_button(3):
#                         print("PP_up")
#                         pygame.event.clear()
#                         break
#             # 方向键（Hat）事件
#             elif event.type == pygame.JOYHATMOTION:
#                 print(f"方向键 {event.hat} 的值: {event.value}")
#                 precise_delay(0.5)